import ast
import cv2
import numpy as np
import onnxruntime as ort



def letterbox(image, new_shape=(640, 640), auto=False, scaleFill=False, scaleUp=True):
    """
    python的信封图片缩放

    总的来说，这个函数的主要目的是将输入的图片缩放到指定的新形状，并根据需要进行填充，以保持图片的纵横比。填充的颜色可以通过 value 参数进行指定，默认为 (114, 114, 114)。函数提供了一些选项，如自动调整填充尺寸、允许向上缩放和进行填充等，以满足不同的需求。

    :param image: 原图
    :param new_shape: 缩放后的图片
    :param color: 填充的颜色
    :param auto: 是否为自动
    :param scaleFill: 填充
    :param scaleUp: 向上填充
    :return:
    """

    # image 的类型是 numpy.ndarray。
    # 也就是说 image 图片数据是一个多维数组，数组中的第一维表示图片的高度，第二维表示宽度，第三维表示像素

    print("预测的图片的类型", type(image), image.shape, image.dtype)

    if isinstance(new_shape, int):
        # 如果 new_shape 是一个整数，则将其转换为 (new_shape, new_shape) 的元组
        new_shape = (new_shape, new_shape)

    target_height = new_shape[0]
    target_width = new_shape[1]

    # 获取原始图片的形状（高度和宽度）
    # 注意，这儿是 image.shape 是一个元组：（image数组的第一维长度，image数组的第二维长度，image数组的第三维长度），（图片的高，图片的宽，像素）
    image_height = image.shape[0]
    image_width = image.shape[1]

    # 宽度和高度的缩放比例
    # ration 成为了元组
    ratio_height = target_height / image_height
    ratio_width = target_width / image_width
    # 计算缩放比例
    radio = min(ratio_height, ratio_width)

    if not scaleUp:
        # 如果不允许向上缩放，则确保缩放比例不超过 1
        radio = min(radio, 1.0)

    # 计算缩放后的尺寸
    new_unpad_width = int(round(image_width * radio))
    new_unpad_height = int(round(image_height * radio))

    # 计算填充的宽度和高度
    dw = target_width - new_unpad_width
    dh = target_height - new_unpad_height

    if auto:
        # 如果 auto 为 True，则将填充尺寸调整为 64 的倍数
        dw = np.mod(dw, 64)
        dh = np.mod(dh, 64)
    elif scaleFill:
        # 如果 scaleFill 为 True，则不进行填充
        # 缩放后的尺寸设置为目标形状
        dw = 0.0
        dh = 0.0
        new_unpad_width = target_width
        new_unpad_height = target_height

    # 均分处理
    # 计算填充的宽度和高度的一半
    dw /= 2
    dh /= 2


    if  image_height != new_unpad_height and image_width != new_unpad_width:
        # 如果原始形状与缩放后的尺寸不同，则使用线性插值进行缩放
        image = cv2.resize(image, (new_unpad_width, new_unpad_height), interpolation=cv2.INTER_LINEAR)

    # 计算顶部和底部的填充尺寸
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    # 计算左侧和右侧的填充尺寸
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))

    # 使用指定的值进行填充
    image = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))  # 添加边界

    meta = {
        "original_height": image_height,
        "original_width": image_width,
        "ratio_height": ratio_height,
        "ratio_width": ratio_width,
        "ratio": radio,
        "padding_top": top,
        "padding_button": bottom,
        "padding_left": left,
        "padding_right": right,
    }

    # 返回缩放和填充后的图片
    return image, meta


def read_image(path):
    # cv2 读取图片得到的数据是 numpy 数组类型
    image_source = cv2.imread(path)

    return image_source


def convert_image(image_source):
    """
    图像预处理，以便在深度学习模型中使用。

    主要的功能是对输入的图像进行一系列处理，包括调整大小、转换颜色通道、归一化和增加维度等。

    :param image_source: 原始图片数据。必须是 NumPy 的数组格式
    :return:
    """
    image_detect = image_source.copy()

    # 图片预处理
    image_detect, meta = letterbox(image_detect, new_shape=640)

    # 反转颜色通道
    #
    # 将图像的 BGR 格式（通常在 OpenCV 中使用）转换为 RGB 格式（常用于大多数深度学习库）。
    #
    # 这部分代码的含义是对图像进行颜色通道的反转。image_detect 是一个 NumPy 数组，通常在图像处理中，图像的形状为 (height, width, channels)，其中 channels 表示颜色通道。
    #
    # 以下是对每个部分的解释：
    # : 表示选择所有行（高度）。
    # : 表示选择所有列（宽度）。
    # ::-1 表示反转最后一个维度（颜色通道）。
    #
    # 通常情况下，图像是以 RGB 格式存储的，而 OpenCV 等一些库可能会以 BGR 格式存储。因此，通过这种方式反转颜色通道，可以将 BGR 转换为 RGB。
    image_detect = image_detect[:, :, ::-1]

    # 维度重排
    #
    # 将图像的维度从 (height, width, channels) 转换为 (channels, height, width)，这是大多数深度学习框架（如 PyTorch）所需的输入格式。
    #
    # 这一部分代码用于改变数组的维度顺序。transpose 方法用于重新排列数组的维度，其参数为一个元组，表示新维度的顺序。这里的 (2, 0, 1) 说明：
    #
    # 2：新维度的第一维对应原数组的第三维（颜色通道）。
    # 0：新维度的第二维对应原数组的第一维（高度）。
    # 1：新维度的第三维对应原数组的第二维（宽度）。
    #
    # 通过这样的转换，图像的维度会变为 (channels, height, width)，这是深度学习框架（如 PyTorch）中常用的图像格式。
    print("Transformed:", type(image_detect), image_detect.shape)
    image_detect = image_detect.transpose(2, 0, 1)
    print("Transformed:", type(image_detect), image_detect.shape)

    # 确保数组的内存布局是连续的
    #
    # NumPy 中的数组可以是“连续的”或“非连续的”。连续的数组在内存中是按顺序存储的，这对于某些操作和效率来说是非常重要的。
    #
    # 当数组在某些操作（例如切片、转置等）之后，可能会变成非连续的。如果数组是非连续的，某些操作（尤其是 C 或 Fortran 后端的库）可能会变得效率低下，或者甚至不支持。
    # np.ascontiguousarray() 会检查 image_detect 是否为连续数组。如果不是，它会返回一个新的连续数组，数据内容与原数组相同。
    # 例如，假设 image_detect 是一个经过切片或转置的数组，可能会导致其内存布局变得不连续。这个函数会将其转化为连续的数组。
    image_detect = np.ascontiguousarray(image_detect)

    # 转换数组的数据类型为 float32
    #
    # 将数组的元素类型从原始类型转换为 float32（32位浮点数）。
    #
    # 在图像处理和深度学习中，许多操作需要浮点数数据类型（例如，进行归一化、计算梯度等）。
    # 如果原始图像数据是整数类型（如 uint8），在进行计算时转换为浮点数是常见的做法。
    # astype() 方法会返回一个新的数组，元素的类型为指定的类型，而不会改变原始数组。
    print("Transformed dtype :", image_detect.dtype)  # float32
    image_detect = image_detect.astype(np.float32)
    print("Transformed dtype :", image_detect.dtype)  # float32

    # 归一化图像数据
    #
    # 将图像像素值从 0-255 的范围缩放到 0.0-1.0 的范围。
    #
    # 在数字图像中，像素值通常是整数，范围为 0 到 255（对于 8 位图像）。这表示图像的颜色强度。
    # 对于机器学习和深度学习模型，尤其是卷积神经网络（CNN），将输入数据归一化为 0 到 1 的范围可以帮助模型更快地收敛并提高准确性。
    # 使用 image_detect /= 255.0 是一个快速的方式，将数组中的每个元素都除以 255.0，从而将像素值转换为浮点数形式，范围在 0 到 1 之间。
    image_detect /= 255.0

    # 增加数组的维度
    #
    # 为 image_detect 数组添加一个新的维度，使其适合于输入到深度学习模型。
    #
    # np.expand_dims() 是一个 NumPy 函数，用于在指定的轴上增加数组的维度。
    # axis=0 表示在数组的第一个维度（最外层维度）上增加一个新的维度。这通常用于将单个图像数据转化为批量图像数据格式（即批处理）。
    # 例如，如果 image_detect 最初的形状是 (height, width, channels)，在执行 np.expand_dims(image_detect, axis=0) 后，
    # 数组的形状将变为 (1, height, width, channels)。
    # 这表示一个批量，包含一张图像。这样处理后图像就适合以批量形式输入到模型中（即一个批次包含一张图像）。
    print("New shape after expanding dimensions:", image_detect.shape)
    image_detect = np.expand_dims(image_detect, axis=0)
    print("New shape after expanding dimensions:", image_detect.shape)

    return image_detect, meta


def onnx_session_load(onnx_path):
    # 加载 ONNX 模型
    session = ort.InferenceSession(onnx_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])

    return session


def onnx_session_run(session, image):
    # 运行模型进行推理，指定输出和输入
    inputs = session.get_inputs()
    outputs = session.get_outputs()

    print("inputs :", type(inputs), len(inputs), inputs[0])
    print("outputs:", type(outputs), len(outputs), outputs[0])

    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    # 初步学习使用 run 来预测图片上的目标
    outputs = session.run([output_name], {input_name: image })

    return outputs


def draw_boxes(targets, image, convert_meta):
    # 置信度的过滤线，阈值
    obj_conf = 0.25

    names_dict = ast.literal_eval(session.get_modelmeta().custom_metadata_map.get("names"))

    for target in targets:
        # 每一个 target 就是识别到的一个目标，一个方框。

        target_box_conf = target[4]

        # 第一次筛选
        if target_box_conf < obj_conf:
            continue

        # 多个类别的置信度中的最大置信度的值
        target_max_score_value = 0
        # 多个类别的置信度中的最大置信度的类别的索引
        target_max_score_index = 0
        # 类别的置信度数据是从 target 数组的索引位置 5 的地方开始
        temp_index = 5

        while temp_index < len(target):

            score = target[temp_index]

            # 寻找最大的置信度值及类别索引
            if score > target_max_score_value:
                target_max_score_value = score
                target_max_score_index = temp_index

            temp_index += 1

        if target_max_score_index < 6:
            continue

        target_max_score_index = target_max_score_index - 5

        # 第二次筛选
        if target_max_score_value * target_box_conf < obj_conf:
            continue

        # box_center_x, box_center_y, box_width, box_height 分别是目标中心的位置和宽高
        box_center_x, box_center_y, box_width, box_height = int(target[0]), int(target[1]), int(target[2]), int(target[3])

        # 将 box_center_x, box_center_y, box_width, box_height 转为 (x1, y1) (x2, y2) 坐标
        # 定义方框的位置（左上角和右下角坐标）
        left_x = box_center_x - (box_width / 2)
        left_y = box_center_y - (box_height / 2)
        right_x = box_center_x + (box_width / 2)
        right_y = box_center_y + (box_height / 2)

        left_x = (left_x - convert_meta["padding_left"]) / convert_meta["ratio"]
        left_y = (left_y - convert_meta["padding_top"]) / convert_meta["ratio"]
        right_x = (right_x - convert_meta["padding_left"]) / convert_meta["ratio"]
        right_y = (right_y - convert_meta["padding_top"]) / convert_meta["ratio"]

        name = names_dict.get(target_max_score_index)

        # 在图像上绘制方框
        cv2.rectangle(image, (int(left_x), int(left_y)), (int(right_x), int(right_y)), (0, 0, 255), thickness=2)
        cv2.putText(image, '{0}--{1:.2f}'.format(name, target_max_score_value), (int(left_x), int(left_y) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), thickness=1)


        print("检测到时目标：", name, (left_x, left_y), (right_x, right_y))

    # 显示图像
    cv2.imshow("detector", image)
    cv2.waitKey(0)



if __name__ == '__main__':
    # 第一步：读取图片文件
    # image_source 是 numpy 数组
    image_source = read_image("D:\\Workspace\\deal\\ada0c42e8bc94f559be0fe0a74644535.jpg")

    print("原始的图片的类型", type(image_source), image_source.shape, image_source.dtype)

    # 第二步：转换图片文件
    # image_detect 是 numpy 数组
    image_detect, convert_meta = convert_image(image_source)

    print("预测的图片的类型", type(image_detect), image_detect.shape, image_detect.dtype)

    # 第三步：加载 ONNX 模型
    session = onnx_session_load("D:\\Workspace\\deal\\yolov5s.onnx")

    # 第四步：检测目标
    # detect_result 是 list 类型
    detect_result = onnx_session_run(session, image_detect)

    print("预测结果：")
    print("detect_result:", type(detect_result), len(detect_result))
    print(detect_result)

    print()

    detect_result = detect_result[0][0]
    print("预测到的目标数量为：", len(detect_result))

    # 第五步：绘制目标
    draw_boxes(detect_result, image_source, convert_meta)
